# 导入所需库
import tushare as ts
import numpy as np
import pandas as pd
import talib
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report
from sklearn.model_selection import train_test_split, GridSearchCV

# 设置 matplotlib 支持中文
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 初始化 tushare pro 接口，需要替换为你自己的 token
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')

try:
    # 获取股票行情数据
    df = pro.daily(ts_code='920118.BJ', start_date='2021', end_date='2025')

    # 数据预处理
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df.set_index('trade_date', inplace=True)

    # 计算技术指标与特征值
    # 计算简单衍生变量
    df['close-open'] = (df['close'] - df['open']) / df['open']
    df['high-low'] = (df['high'] - df['low']) / df['low']
    df['pre_close'] = df['close'].shift(1)
    df['price_change'] = df['close'] - df['pre_close']
    df['p_change'] = (df['close'] - df['pre_close']) / df['pre_close'] * 100

    # 计算移动平均线指标
    df['MA5'] = df['close'].rolling(5).mean()
    df['MA10'] = df['close'].rolling(10).mean()

    # 用 TA-Lib 库生成指标
    df['RSI'] = talib.RSI(df['close'].values, timeperiod=12)
    df['MOM'] = talib.MOM(df['close'].values, timeperiod=5)
    df['EMA12'] = talib.EMA(df['close'].values, timeperiod=12)
    df['EMA26'] = talib.EMA(df['close'].values, timeperiod=26)
    df['MACD'], df['MACDsignal'], df['MACDhist'] = talib.MACD(df['close'].values, fastperiod=6, slowperiod=12,
                                                              signalperiod=9)

    # 去除包含缺失值的行
    df.dropna(inplace=True)

    # 特征和目标变量
    features = df[['close', 'vol', 'close-open', 'MA5', 'MA10', 'high-low', 'RSI', 'MOM', 'EMA12', 'MACD', 'MACDsignal',
                   'MACDhist']]
    target = np.where(df['price_change'].shift(-1) > 0, 1, -1)

    # 去除最后一行，因为目标变量最后一行是 NaN
    features = features[:-1]
    target = target[:-1]

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(features, target, test_size=0.2, random_state=42)

    # 建立随机森林模型
    model = RandomForestClassifier(random_state=42)

    # 参数调优
    param_grid = {
        'n_estimators': [50, 100, 200],
        'max_depth': [5, 10, 15],
        'min_samples_split': [2, 5, 10],
        'min_samples_leaf': [1, 2, 4]
    }

    grid_search = GridSearchCV(model, param_grid, cv=5, scoring='accuracy')
    grid_search.fit(X_train, y_train)

    # 最优模型
    best_model = grid_search.best_estimator_

    # 预测
    y_pred = best_model.predict(X_test)

    # 模型评价
    accuracy = accuracy_score(y_test, y_pred)
    print(f"模型准确率: {accuracy}")
    print("分类报告:")
    print(classification_report(y_test, y_pred))

    # 在测试数据上添加一列，预测收益
    X_test['prediction'] = best_model.predict(X_test)

    # 计算每天的股价变化率
    X_test['p_change'] = (X_test['close'] - X_test['close'].shift(1)) / X_test['close'].shift(1)
    X_test['p_change'] = X_test['p_change'].fillna(0)  # 处理缺失值

    # 计算累积收益率
    X_test['origin'] = (X_test['p_change'] + 1).cumprod()

    # 计算利用模型预测后的收益率
    X_test['strategy'] = (X_test['prediction'].shift(1) * X_test['p_change'] + 1).cumprod()

    # 绘制收益曲线
    plt.figure(figsize=(12, 6))
    X_test[['strategy', 'origin']].dropna().plot()
    plt.title('920118.BJ 策略收益率 vs 简单持有收益率')
    plt.xlabel('日期')
    plt.ylabel('累积收益率')
    plt.legend(['策略收益率', '简单持有收益率'])
    plt.gcf().autofmt_xdate()
    plt.show()

except Exception as e:
    print(f"发生错误: {e}")
